import os
import sys
import numpy as np

from OCC.Extend.DataExchange import (
    read_iges_file,
    read_step_file,
    read_stl_file,
    write_iges_file,
    write_step_file,
    write_stl_file,
)

from OCC.Core.BRepPrimAPI import (
    BRepPrimAPI_MakeSphere,
    BRepPrimAPI_MakeTorus,
    BRepPrimAPI_MakeBox,
    BRepPrimAPI_MakeCone,
    BRepPrimAPI_MakeCylinder,
    BRepPrimAPI_MakeWedge,
    BRepPrimAPI_MakePrism,
)

from OCC.Core.BRepAlgoAPI import (
    BRepAlgoAPI_Fuse,
    BRepAlgoAPI_Cut,
    BRepAlgoAPI_Common,
)
from OCC.Core.BRepBuilderAPI import (
    BRepBuilderAPI_MakeVertex,
    BRepBuilderAPI_MakeEdge,
    BRepBuilderAPI_MakeWire,
    BRepBuilderAPI_MakeFace,
)
from OCC.Core.BRep import BRep_Tool
from OCC.Core.GC import (
    GC_MakeSegment,
    GC_MakeArcOfCircle,
)
from OCC.Core.GCE2d import GCE2d_MakeSegment
from OCC.Core.GeomAPI import geomapi_To3d

from OCC.Core.gp import (
    gp_Pnt,
    gp_Pnt2d,
    gp_Vec,
    gp_Pln,
    gp_Ax2,
    gp_Dir,
    gp_Circ,
    gp_Elips,
)

from OCC.Core.TopAbs import (
    TopAbs_SOLID,
    TopAbs_EDGE,
    TopAbs_FACE,
)
from OCC.Core.TopoDS import (
    TopoDS_Compound,
    TopoDS_Vertex,
    TopoDS_Wire,
    TopoDS_Face
)

from OCC.Core.Bnd import Bnd_Box
from OCC.Core.BRepBndLib import brepbndlib_Add
from OCC.Core.BRepMesh import BRepMesh_IncrementalMesh

import meshio
import gmsh
import pygmsh

from traits.api import *
from traitsui.api import *

import tempfile


class TraitsClass(HasTraits):
    shape = None
    mesh = None
    v = None
    f = None

    def generateBox(self):
        self.shape = BRepPrimAPI_MakeBox(10., 20., 30.).Shape()

    def importStep(self, filename):
        self.shape = read_step_file(filename)

    def importIges(self, filename):
        self.shape = read_iges_file(filename)

    def importStl(self, filename):
        self.mesh = read_stl_file(filename)

    def exportStep(self, filename):
        write_step_file(self.shape, filename)

    def exportIges(self, filename):
        write_iges_file(self.shape, filename)

    def importMesh(self, filename):
        mesh = meshio.read(filename)
        self.v = mesh.points
        self.f = mesh.cells_dict['triangle']
        self.meshToShape()

    def meshToShape(self):
        if self.v is None or self.f is None:
            return
        file = tempfile.NamedTemporaryFile(suffix='.stl')
        file.close()
        meshio.write_points_cells(file.name, self.v, [('triangle', self.f)])
        self.mesh = read_stl_file(file.name)

    def shapeToMesh(self):
        if not self.shape:
            return
        gmsh.initialize()
        gmsh.clear()
        file = tempfile.NamedTemporaryFile(suffix='.stp')
        file.close()
        write_step_file(self.shape, file.name)
        gmsh.model.occ.importShapes(file.name)
        gmsh.model.occ.synchronize()
        gmsh.option.setNumber("Mesh.MeshSizeMin", 500)
        gmsh.option.setNumber("Mesh.MeshSizeMax", 500)
        gmsh.model.mesh.generate(2)
        file = tempfile.NamedTemporaryFile(suffix='.stl')
        file.close()
        gmsh.write(file.name)
        self.importMesh(file.name)

    def shapeSize(self):
        if self.shape is None:
            return
        bbox = Bnd_Box()
        brepbndlib_Add(self.shape, bbox)
        xmin, ymin, zmin, xmax, ymax, zmax = bbox.Get()
        print(xmin, ymin, zmin, xmax, ymax, zmax, xmax - xmin, ymax - ymin, zmax - zmin)
        return xmax - xmin, ymax - ymin, zmax - zmin

    def meshSize(self):
        if self.mesh is None:
            return
        bbox = Bnd_Box()
        brepbndlib_Add(self.mesh, bbox)
        xmin, ymin, zmin, xmax, ymax, zmax = bbox.Get()
        print(xmin, ymin, zmin, xmax, ymax, zmax, xmax - xmin, ymax - ymin, zmax - zmin)
        return xmax - xmin, ymax - ymin, zmax - zmin

    def traits_init(self):
        pass
        # self.importIges('cube.igs')
        # self.importStep('11752.stp')
        self.generateBox()


if __name__ == "__main__":

    from PyQt5.QtCore import *
    from PyQt5.QtGui import *
    from PyQt5.QtWidgets import *

    from OCC.Display.backend import load_backend

    load_backend("qt-pyqt5")
    import OCC.Display.qtDisplay as qtDisplay


    class MainWindow(QMainWindow):
        a = 4234.42342

        def __init__(self):
            super(MainWindow, self).__init__()
            self.model = TraitsClass()

            self.resize(800, 600)
            menubar = QMenuBar()
            self.setMenuBar(menubar)
            self.statusbar = QStatusBar(self)
            self.setStatusBar(self.statusbar)

            menu = QMenu('模型', menubar)
            menubar.addMenu(menu)

            action = QAction('立方块', self)
            menu.addAction(action)
            action.triggered.connect(lambda: (
                self.model.generateBox(),
                self.displayShape(),
            ))

            action = QAction('导入stp', self)
            menu.addAction(action)
            action.triggered.connect(lambda: (
                self.model.importStep('tmp.stp'),
                self.displayShape(),
            ))

            action = QAction('导入igs', self)
            menu.addAction(action)
            action.triggered.connect(lambda: (
                self.model.importIges('cube.igs'),
                self.displayShape(),
            ))

            def showShapeSize():
                x, y, z = self.model.shapeSize()
                QMessageBox.information(self, '尺寸', f'x轴：{x}\ny轴：{y}\nz轴：{z}')

            action = QAction('尺寸', self)
            menu.addAction(action)
            action.triggered.connect(lambda: (
                showShapeSize()
            ))

            menu = QMenu('网格', menubar)
            menubar.addMenu(menu)

            action = QAction('网格剖分', self)
            menu.addAction(action)
            action.triggered.connect(lambda: (
                self.model.shapeToMesh(),
                self.displayMesh(),
            ))

            action = QAction('导入stl', self)
            menu.addAction(action)
            action.triggered.connect(lambda: (
                self.model.importMesh('delete.stl'),
                self.displayMesh(),
            ))

            def showMeshSize():
                x, y, z = self.model.meshSize()
                QMessageBox.information(self, '尺寸', f'x轴：{x}\ny轴：{y}\nz轴：{z}')

            action = QAction('尺寸', self)
            menu.addAction(action)
            action.triggered.connect(lambda: (
                showMeshSize()
            ))

            menu = QMenu('窗口', menubar)
            menubar.addMenu(menu)

            action = QAction('模型', self)
            menu.addAction(action)
            action.triggered.connect(lambda: (
                self.centralWidget().setParent(None),
                self.setCentralWidget(self.centralwidget),
            ))

            action = QAction('网格', self)
            menu.addAction(action)
            action.triggered.connect(lambda: (
                self.centralWidget().setParent(None),
                self.setCentralWidget(self.centralwidget1),
            ))

            self.centralwidget = QWidget()
            horizontalLayout = QHBoxLayout(self.centralwidget)

            self.canvas = qtDisplay.qtViewer3d(self)
            self.canvas.InitDriver()  # canva的驱动,设置驱动后，才能成功display
            self.display = self.canvas._display
            self.display.set_bg_gradient_color([0, 128, 0], [128, 128, 128])  # 设置背景渐变色
            self.display.display_triedron()  # display black trihedron

            horizontalLayout.addWidget(self.canvas)

            self.centralwidget1 = QWidget()
            horizontalLayout = QHBoxLayout(self.centralwidget1)
            pushButton = QPushButton('haha1')
            horizontalLayout.addWidget(pushButton)
            pushButton.clicked.connect(lambda: (
                print('fjdksf'),
            ))
            pushButton = QPushButton('haha2')
            horizontalLayout.addWidget(pushButton)
            pushButton.clicked.connect(lambda: (
                print('fjdksf12'),
            ))

            pushButton = QPushButton('haha3')
            horizontalLayout.addWidget(pushButton)
            pushButton.clicked.connect(lambda: (
                print('fjdksf13'),
            ))

            self.setCentralWidget(self.centralwidget)

            self.displayShape()

        def displayShape(self):
            if self.model.shape is None:
                return
            self.display.EraseAll()
            self.display.DisplayShape(self.model.shape, update=True)
            self.display.FitAll()

        def displayMesh(self):
            if self.model.mesh is None:
                return
            self.display.EraseAll()
            self.display.DisplayShape(self.model.mesh, update=True)
            self.display.FitAll()


    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())
